
#include "../iface/LibSSH2Interface.h"
#include "../aprtools/APRFile.h"
#include "../aprtools/APRAutoLock.h"
#include "ShellChannel.h"
#include "SCPSendChannel.h"
#include "RemoteSSHKeyFile.h"
#include "../crypt/RSAPublicKey.h"
#include "../crypt/DSAPublicKey.h"
#include "../crypt/ECDSAPublicKey.h"
#include <apr_strings.h>

/** Get the home directory location on the remote server.
 * 
 * @return A string containing the home directory location, to be
 * freed by the caller. */
char *RemoteSSHKeyFile::getHomeDirectory()
{
   ShellChannel tshell(*tclient);
   APRFuture<char *> *rstringfuture = tshell.runCommandToStringFuture(
      "sh -c 'echo -n \"$HOME\"'");
   if (rstringfuture == NULL) return NULL;
   char *result_string = *(rstringfuture->get());
   delete rstringfuture;
   return result_string;
}

/** Checks whether the remote file given is writable only. If it does not
 * exist, this will fail. Also works on directories.
 * 
 * @param filename The name of the file to query for writability.
 * 
 * @return One if it is writable, zero if it is not writable. */
int RemoteSSHKeyFile::checkRemoteFileWritableOnly(const char *filename)
{
   ShellChannel tshell(*tclient);
   APRFuture<char *> *rstringfuture = tshell.runCommandToStringFuture(
      "sh -c 'test -w \"%s\"; echo $?'", filename);
   if (rstringfuture == NULL) return 0;
   char *result_string = *(rstringfuture->get());
   int tretval = !(atoi(result_string));
   free(result_string);
   delete rstringfuture;
   return tretval;
}

/** Checks whether the remote file given is writable. If it does not
 * exist, this will fail. Also works on directories.
 * 
 * @param filename The name of the file to query for writability.
 * 
 * @return One if it is writable, zero if it is not writable. */
int RemoteSSHKeyFile::checkRemoteFileWritable(const char *filename)
{
   // Try touching it first.
   ShellChannel tshelltouch(*tclient);
   APRFuture<char *> *rstringtouchfuture = tshelltouch.runCommandToStringFuture(
      "sh -c 'mkdir -p \"$(dirname \"%s\")\"; touch \"%s\"'", filename, filename);
   if (rstringtouchfuture == NULL) return 0;
   char *result_touch_string = *(rstringtouchfuture->get());
   free(result_touch_string);
   delete rstringtouchfuture;
   ShellChannel tshell(*tclient);
   APRFuture<char *> *rstringfuture = tshell.runCommandToStringFuture(
      "sh -c 'test -w \"%s\"; echo $?'", filename);
   if (rstringfuture == NULL) return 0;
   char *result_string = *(rstringfuture->get());
   int tretval = !(atoi(result_string));
   free(result_string);
   delete rstringfuture;
   return tretval;
}

/** Tries to obtain a lock on the remote file by turning on 'no
 * clobber' and writing to a remote file. It is atomic because two
 * files cannot clash, one cannot fail quietly and the error condition
 * is detectable and a lock failure.
 * 
 * @param lockdir The directory in which to create the lock.
 * 
 * @param lockname The name of the file in which to create the lock.
 * 
 * @return Greater than zero specifies the PID of the locking process,
 * less than zero specifies an error. */
int RemoteSSHKeyFile::tryLock(const char *lockdir, const char *lockname)
{
   ShellChannel tshell(*tclient);
   APRFuture<char *> *rstringfuture = tshell.runCommandToStringFuture(
      "sh -c 'set -C; if echo \"$$\" > %s/%s; then echo \"$$\"; else echo -1; fi'",
      lockdir, lockname);
   if (rstringfuture == NULL) return -1;
   char *result_string = *(rstringfuture->get());
   int tretval = atoi(result_string);
   free(result_string);
   delete rstringfuture;
   return tretval;
}

/** Releases a remote lock with the matching PID number.
 * 
 * @param lockdir The directory in which to create the lock.
 * 
 * @param lockname The name of the file in which to create the lock.
 * 
 * @param pidnum The number of the PID which was recorded in the
 * locking phase.  This is compared against the PID recorded in the
 * file, and the lock is released successfully only when they match.
 * 
 * @return Zero on success, non-zero (-1) on failure. */
int RemoteSSHKeyFile::releaseMatchingLock(const char *lockdir, const char *lockname, int pidnum)
{
   ShellChannel tshell(*tclient);
   APRFuture<char *> *rstringfuture = tshell.runCommandToStringFuture(
      "sh -c 'set -C; if [ `cat %s/%s` -eq %d ]; then rm %s/%s; echo 0; else echo -1; fi'",
      lockdir, lockname, pidnum, lockdir, lockname);
   if (rstringfuture == NULL) return -1;
   char *result_string = *(rstringfuture->get());
   int tretval = atoi(result_string);
   free(result_string);
   delete rstringfuture;
   return tretval;
}

/** Forcibly release a remote lock, regardlesss of what PID number is
 * contained.
 * 
 * @param lockdir The directory in which to create the lock.
 * 
 * @param lockname The name of the file in which to create the
 * lock. */
void RemoteSSHKeyFile::forceRemoveLockFile(const char *lockdir, const char *lockname)
{
   errprint("forcibly removing stale remote lockfile '%s/%s'", lockdir, lockname);
   ShellChannel tshell(*tclient);
   APRFuture<char *> *rstringfuture = tshell.runCommandToStringFuture(
      "sh -c 'rm %s/%s'",
      lockdir, lockname);
   char *result_string = *(rstringfuture->get());
   if (result_string) free(result_string);
   delete rstringfuture;
}

/** Dump a remote file into a string. Useful for debugging.
 * 
 * @param filename Remote file name to dump.
 * 
 * @return String of th contents of the remote file, to be freed by
 * the caller. */
char *RemoteSSHKeyFile::dumpFileData(const char *filename)
{
   ShellChannel tshell(*tclient);
   APRFuture<char *> *rstringfuture = tshell.runCommandToStringFuture(
      "sh -c 'cat \"%s\"'", filename);
   if (rstringfuture == NULL) return NULL;
   char *result_string = *(rstringfuture->get());
   delete rstringfuture;
   return result_string;
}

/** Report all keys. Give a report of all the managed keys present. */
void RemoteSSHKeyFile::reportAllKeys()
{
   for (int i = 0; i < file_keytypes.size(); i++)
   {
      const char *ktype = "UNKNOWN";
      switch (file_keytypes.at(i))
      {
         case SSHKEY_RSA:
            ktype = "RSA"; break;
         case SSHKEY_DSA:
            ktype = "DSA"; break;
         case SSHKEY_ECDSA521:
            ktype = "ECDSAp521"; break;
         case SSHKEY_ECDSA384:
            ktype = "ECDSAp384"; break;
         case SSHKEY_ECDSA256:
            ktype = "ECDSAp256"; break;
         default: break;
      }
      if (file_ismanaged.at(i))
         errprint("REMOTE KEY FOUND: %s %s %s", ktype,
            (file_ismanaged.at(i)) ? "MANAGED" : "UNMANAGED", (file_keytag.at(i)));
   }
}

/** Import a key from a line in the key file data buffer.
 * 
 * @param keyfiledata The dumped data from the key file.
 * 
 * @param beginidx The beginning index of the substring of the file
 * data containing the key.
 * 
 * @param endidx The end index of the substring of the file data
 * containing the key. */
void RemoteSSHKeyFile::importKeyFromLine(const char *keyfiledata, int beginidx, int endidx)
{
   char *tpstring = (char *)malloc(sizeof(char)*(1 + endidx - beginidx));
   memcpy(tpstring, &(keyfiledata[beginidx]), sizeof(char)*(endidx - beginidx));
   tpstring[endidx - beginidx] = '\0';
   char *first_token, *second_token;
   char *next_token;
   first_token = apr_strtok(tpstring, " ", &next_token);
   second_token = apr_strtok(NULL, " ", &next_token);
   if ((first_token == NULL) ||
       (second_token == NULL) ||
       (next_token == NULL))
   {
      free(tpstring);
      return;
   }
   sshkeytype t_keytype = SSHKEY_UNKNOWN;
   if (strcmp(first_token, "ssh-rsa") == 0)
      t_keytype = SSHKEY_RSA;
   if (strcmp(first_token, "ssh-dss") == 0)
      t_keytype = SSHKEY_DSA;
   if (strcmp(first_token, "ecdsa-sha2-nistp521") == 0)
      t_keytype = SSHKEY_ECDSA521;
   if (strcmp(first_token, "ecdsa-sha2-nistp384") == 0)
      t_keytype = SSHKEY_ECDSA384;
   if (strcmp(first_token, "ecdsa-sha2-nistp256") == 0)
      t_keytype = SSHKEY_ECDSA256;
   char *dissectable = strdup(next_token);
   char *t_keytag = NULL;
   char *try_our_tag = apr_strtok(dissectable, "_", &t_keytag);
   int t_managed = 0;
   //printf(" %s %s\n", try_our_tag, ourtag);
   if (strcmp(try_our_tag, ourtag) == 0)
   {
      t_managed = 1;
   } else
   {
      t_managed = 0;
      t_keytag = next_token;
   }
   file_keytypes.push_back(t_keytype);
   file_keytypetag.push_back(strdup(first_token));
   file_ismanaged.push_back(t_managed);
   // should probably check whether this key tag is already in the class?
   file_keytag.push_back(strdup(t_keytag));
   file_keybase64.push_back(strdup(second_token));
   free(dissectable);
   free(tpstring);
}

/** Synchronise the class with the remote key file.
 * 
 * @param keyfilename The path and file name of the remote file
 * containing the users public SSH keys for identification, usually
 * "~/.ssh/authorized_keys". */
void RemoteSSHKeyFile::remoteKeySync(const char *keyfilename)
{
   char *keyfiledata = dumpFileData(keyfilename);
   int filelength = strlen(keyfiledata);
   char *buffer = (char *)malloc(sizeof(char)*(filelength+1));
   int beginidx = 0;
   int endidx = 0;
   while ((beginidx < filelength) && (endidx < filelength))
   {
      if (keyfiledata[endidx] != '\n')
         endidx++;
      else
      {
         importKeyFromLine(keyfiledata, beginidx, endidx);
         endidx++;
         beginidx = endidx;
      }
   }
   free(buffer);
   free(keyfiledata);
   return;
}

/** Commit key changes by updating the remote file to reflect the
 * (managed only) public SSH keys present in the local class.
 * 
 * @param keyfilename Remote key file name, usually
 * "~/.ssh/authorized_keys".
 * 
 * @param keytempdir A temporary directory to write the modified key
 * file into.
 * 
 * @param keytempfile The file within the temporary directory to write
 * the modified public SSH key data for this logged in user into.
 * 
 * @return Zero on success, non-zero on failure. */
int RemoteSSHKeyFile::commitKeyChanges(const char *keyfilename, const char *keytempdir, const char *keytempfile)
{
   // serialise and write
   APRBucketBrigade aprbb(imp);
   for (int i = 0; i < file_keytypes.size(); i++)
   {
      aprbb.writefToStream("%s %s %s%s%s\n", file_keytypetag.at(i), file_keybase64.at(i),
         (file_ismanaged.at(i)) ? ourtag : "", (file_ismanaged.at(i)) ? "_" : "", file_keytag.at(i));
   }
   // check that both files can be written to
   if (!checkRemoteFileWritable(keytempdir))
   {
      errprint("could not write to remote temp directory");
      return 1;
   }
   if (!checkRemoteFileWritable(keyfilename))
   {
      errprint("could not write to main key file");
      return 1;
   }
   // generate SCP to file destination
   char *remote_temp_file_name = StringUtils::varprintf("%s/%s", keytempdir, keytempfile);
   char *newkeyfiledata = NULL;
   size_t buffer_length = 0;
   aprbb.writeEOS();
   aprbb.bucketBrigadeToBuffer(&newkeyfiledata, &buffer_length);
   SCPSendChannel *scpsendchan = new SCPSendChannel(*tclient);
   //printf("%s %d\n", newkeyfiledata, buffer_length);
   APRFuture<int> *copyfuture = scpsendchan->copyBufferToRemoteFileFuture(remote_temp_file_name, newkeyfiledata, buffer_length, 0700);
   if (copyfuture == NULL) return 1;
   int result = *(copyfuture->get());
   delete copyfuture;
   delete scpsendchan;
   free(newkeyfiledata);
   if (result != 0)
   {
      errprint("something went wrong with the secure copy");
      free(remote_temp_file_name);
      return result;
   }
   ShellChannel tshell2(*tclient);
   APRFuture<char *> *rstringfuture = tshell2.runCommandToStringFuture(
      "sh -c 'mv -f \"%s\" \"%s\"'", remote_temp_file_name, keyfilename);
   if (rstringfuture == NULL) return 1;
   char *result_string = *(rstringfuture->get());
   delete rstringfuture;
   free(result_string);
   if (checkRemoteFileWritableOnly(remote_temp_file_name))
   {
      errprint("remote temp file is writable - must exist so move failed!");
      ShellChannel tshell(*tclient);
      APRFuture<char *> *rstringfuture = tshell.runCommandToStringFuture(
         "sh -c 'rm -f \"%s\"'", remote_temp_file_name);
      if (rstringfuture == NULL) return 1;
      char *result_string = *(rstringfuture->get());
      free(result_string);
      delete rstringfuture;
      free(remote_temp_file_name);
      return 1;
   }
   // probably okay if it reaches here...
   free(remote_temp_file_name);
   return 0;
}

/** Clean keys by removing any remote keys which have been depracated
 * locally and are no longer in use. */
void RemoteSSHKeyFile::cleanKeys()
{
   // remote keys that are marked deleting (then decrement in the
   // database) or do not exist
   std::vector<sshkeytype> file_copy_keytypes;
   std::vector<int> file_copy_ismanaged;
   std::vector<char *> file_copy_keytypetag;
   std::vector<char *> file_copy_keytag;
   std::vector<char *> file_copy_keybase64;
   // copy to temporary holding
   for (int i = 0; i < file_keytypes.size(); i++)
   {
      file_copy_keytypes.push_back(file_keytypes.at(i));
      file_copy_keytypetag.push_back(file_keytypetag.at(i));
      file_copy_ismanaged.push_back(file_ismanaged.at(i));
      file_copy_keytag.push_back(file_keytag.at(i));
      file_copy_keybase64.push_back(file_keybase64.at(i));
   }
   file_keytypetag.clear();
   file_keytypes.clear();
   file_keytag.clear();
   file_keybase64.clear();
   file_ismanaged.clear();
   // for each key check...
   for (int i = 0; i < file_copy_keytypes.size(); i++)
   {
      int copy_key = 1;
      if (file_copy_ismanaged.at(i))
      {
         const KeyTableView *keytableviews[3] = { NULL, NULL, NULL };
         keytableviews[0] = keydb->signingTable();
         keytableviews[1] = keydb->checkingTable();
         int count = 0;
         while (keytableviews[count] != NULL)
         {
            // check if the key exists in the database table
            if (keytableviews[count]->testKeyExists(file_copy_keytag.at(i)) == 0)
            {
               // if key doesn't exist, and it is tagged as managed by
               // us, then we should delete it!
               errprint("not keeping key '%s' remotely, does not exist in table view (%d)", file_copy_keytag.at(i), count);
               copy_key = 0;
            } else
            {
               // key exists, so we should check if it is marked as
               // deprecated!
               if (keytableviews[count]->getKeyUsedReferenceCount(file_copy_keytag.at(i)) == 0)
               {
                  // if the key is now unused, then
                  errprint("not keeping key '%s' remotely, has zero uses in table view (%d)", file_copy_keytag.at(i), count);
                  copy_key = 0;
                  keytableviews[count]->decrementKeyExistsReferenceCount(file_copy_keytag.at(i));
                  // we can't actually delete the key here, doing so
                  // creates a race condition with replacement keys if
                  // (keytableviews[count]->getKeyExistsReferenceCount(file_copy_keytag.at(i))
                  // == 0)
                  // keytableviews[count]->deleteKey(file_copy_keytag.at(i));
               }
            }
            count++;
         }
      }
      if (copy_key)
      {
         file_keytypes.push_back(file_copy_keytypes.at(i));
         file_keytypetag.push_back(file_copy_keytypetag.at(i));
         file_keytag.push_back(file_copy_keytag.at(i));
         file_keybase64.push_back(file_copy_keybase64.at(i));
         file_ismanaged.push_back(file_copy_ismanaged.at(i));
      } else
      {
         // don't copy, but free the strings!
         free(file_copy_keybase64.at(i));
         free(file_copy_keytag.at(i));
         free(file_copy_keytypetag.at(i));
      }
   }
}

/** Get the base 64 SSH public key format data from a local key, by
 * specifying the key name. This is unique as the only keys permitted
 * in the file are checking keys.
 * 
 * @param keyname Name of the key in the checking key table to
 * transform into SSH format.
 * 
 * @param keypass The password to unlock the encrypted key from the
 * database.
 * 
 * @return The public key base 64 data in SSH format, ready for
 * incorporation into a line of the authorized_keys file. */
char *RemoteSSHKeyFile::getBase64Blob(const char *keyname, const char *keypass)
{
   if ((keyname == NULL) || (strlen(keyname) == 0))
   {
      errprint("key name cannot be empty or NULL");
      return NULL;
   }
   if (keydb == NULL)
   {
      errprint("no database - no key '%s' to add", keyname);
      return NULL;
   }
   CheckingKey *chk_key = Key::getCheckingKey(keydb, keypass, keyname);
   char *keydata = NULL;
   switch (chk_key->getKeyType())
   {
      case AbstractKey::KEY_TYPE_PUBLIC_RSA:
         dynamic_cast<RSAPublicKey *>(chk_key)->getSSHKeyBlob(&keydata, keyname); break;
      case AbstractKey::KEY_TYPE_PUBLIC_DSA:
         dynamic_cast<DSAPublicKey *>(chk_key)->getSSHKeyBlob(&keydata, keyname); break;
      case AbstractKey::KEY_TYPE_PUBLIC_ECDSA:
         dynamic_cast<ECDSAPublicKey *>(chk_key)->getSSHKeyBlob(&keydata, keyname); break;
      default: break;
   }
   delete chk_key;
   if (keydata == NULL)
   {
      errprint("can't produce SSH key blob");
      return NULL;
   }
   // hackily snip out the base64 from within - SSHKeyBlob seems
   // over-zealous
   char *remainder = NULL;
   // it's the second block, as the first will always contain the data
   char *tok1 = apr_strtok(keydata, " ", &remainder);
   char *tok2 = apr_strtok(NULL, " ", &remainder);
   char *result = strdup(tok2);
   free(keydata);
   return result;
}

/** Internally add key, without checks for usefulness.
 * 
 * @param keyname The key name to add to the list of public SSH keys
 * in the local copy of the authorized_keys.
 * 
 * @param keypass The key password to use to get the corresponding
 * checking key from the key database.
 * 
 * @return Zero on success, non-zero on failure. */
int RemoteSSHKeyFile::addKey_int(const char *keyname, const char *keypass)
{
   if ((keyname == NULL) || (strlen(keyname) == 0))
   {
      errprint("key name cannot be empty or NULL");
      return -1;
   }
   if (keydb == NULL)
   {
      errprint("no database - no key '%s' to add", keyname);
      return -1;
   }
   CheckingKey *chk_key = Key::getCheckingKey(keydb, keypass, keyname);
   sshkeytype k_type;
   char *keydata = getBase64Blob(keyname, keypass);
   const char *keytypetag = NULL;
   switch (chk_key->getKeyType())
   {
      case AbstractKey::KEY_TYPE_PUBLIC_RSA:
      {
         k_type = SSHKEY_RSA;
         keytypetag = "ssh-rsa";
         break;
      }
      case AbstractKey::KEY_TYPE_PUBLIC_DSA:
      {
         k_type = SSHKEY_DSA;
         keytypetag = "ssh-dss";
         break;
      }
      case AbstractKey::KEY_TYPE_PUBLIC_ECDSA:
      {
         switch (ECDSA_BITS)
         {
            case 521:
               k_type = SSHKEY_ECDSA521;
               keytypetag = "ecdsa-sha2-nistp521";
            case 384:
               k_type = SSHKEY_ECDSA384;
               keytypetag = "ecdsa-sha2-nistp384";
            case 256:
               k_type = SSHKEY_ECDSA256;
               keytypetag = "ecdsa-sha2-nistp256";
            default: break;
         }
         break;
      }
      default: break;
   }
   if (keydata == NULL)
   {
      errprint("can't produce SSH key blob");
      delete chk_key;
      return 1;
   }
   file_keytypes.push_back(k_type);
   file_keytypetag.push_back(strdup(keytypetag));
   file_ismanaged.push_back(1);
   file_keytag.push_back(strdup(keyname));
   file_keybase64.push_back(keydata);
   delete chk_key;
   return 0;
}

/** Add key to the list if it is not there, incrementing the database
 * use count either way.
 * 
 * @param keyname The key name to add to the list of public SSH keys
 * in the local copy of the authorized_keys.
 * 
 * @param keypass The key password to use to get the corresponding
 * checking key from the key database.
 * 
 * @return Zero on success, non-zero on failure. */
int RemoteSSHKeyFile::addKey(const char *keyname, const char *keypass)
{
   // test if the key is already in the remote file
   int key_exists_flag = 0;
   for (int i = 0; i < file_keytag.size(); i++)
   {
      if (strcmp(file_keytag.at(i), keyname) == 0)
      {
         // key exists
         key_exists_flag = 1;
      }
   }
   // if it is not, then increment the exists counters
   if (key_exists_flag == 0)
   {
      if (addKey_int(keyname, keypass) == 0)
      {
         if (keydb->signingTable()->testKeyExists(keyname))
            keydb->signingTable()->incrementKeyExistsReferenceCount(keyname);
         if (keydb->checkingTable()->testKeyExists(keyname))
            keydb->checkingTable()->incrementKeyExistsReferenceCount(keyname);
      } else
      {
         errprint("error adding key");
      }
   }
}

/** Constructor for the remote SSH key file, synchronising keys in the
 * process.
 * 
 * @param itclient The SSH client over which whose keys are
 * synchronised with.
 * 
 * @param inetdb The network database, needed if the keys require
 * updating.
 * 
 * @param ikeydb The key database for loading relevant keys from.
 * 
 * @param iourtag The string tag in the authorized_keys file that
 * tells us that the key is managed by us. */
RemoteSSHKeyFile::RemoteSSHKeyFile(SSHClient &itclient, NetworkDB &inetdb,
   KeyDB &ikeydb, const char *iourtag) : SSHChannel(itclient)
{
   netdb = &inetdb;
   keydb = &ikeydb;
   ourtag = (iourtag) ? strdup(iourtag) : NULL;
   char *homedirname = getHomeDirectory();
   if (homedirname == NULL)
   {
      errprint("failed to get home directory");
      sshkeyfile = NULL;
      return;
   }
   sshkeyfile = StringUtils::varprintf("%s/.ssh/authorized_keys", homedirname);
   remoteKeySync(sshkeyfile);
   int sleep_amt = 50000;
   int tries = 3;
   free(homedirname);
   while ((lockpid = tryLock("/tmp", "authorized_keys.lock")) <= 0)
   {
      apr_sleep(sleep_amt);
      sleep_amt = (sleep_amt < 500000) ? (sleep_amt << 1) : 500000;
      tries--;
      if (tries == 0) break;
   }
   if (tries == 0)
   {
      forceRemoveLockFile("/tmp", "authorized_keys.lock");
      lockpid = tryLock("/tmp", "authorized_keys.lock");
   }
   reportAllKeys();
   cleanKeys();
}

/** Destructor for the remote SSH key file, saving the new keyfile to
 * the remote server in the process. */
RemoteSSHKeyFile::~RemoteSSHKeyFile()
{
//   reportAllKeys();
   if (sshkeyfile)
      commitKeyChanges(sshkeyfile, "/tmp", "authorized_keys.written");
   if (lockpid)
   {
      if (releaseMatchingLock("/tmp", "authorized_keys.lock", lockpid) != 0)
      {
         errprint("release matching lock failed (remote lock pid '%d')", lockpid);
      }
   } else forceRemoveLockFile("/tmp", "authorized_keys.lock");
   // clean up all the strings in the vectors
   for (int i = 0; i < file_keytypes.size(); i++)
   {
      free(file_keytypetag.at(i));
      free(file_keytag.at(i));
      free(file_keybase64.at(i));
   }
   // need a commit key changes to remote file function here!
   if (ourtag) free(ourtag);
   ourtag = NULL;
   if (sshkeyfile) free(sshkeyfile);
   sshkeyfile = NULL;
}
